# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：model.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/10/11 10:24
@Note: Something beautiful is about to happen !
'''
from dataclasses import dataclass, field
from typing import Any, Tuple, List

import keras.src.callbacks
import matplotlib.pyplot as plt
from keras.src.datasets import mnist
from keras import utils, Model, Sequential, optimizers, losses, metrics
from keras.src.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout

import numpy as np

# 1.	模型约束描述
# 1)	数据预处理：
# 加载MNIST数据集，并将其分为训练集和测试集。
# 从MNIST数据集中精选以下4个类别：
# 数字"1"（类别编号1）
# 数字"3"（类别编号3）
# 数字"7"（类别编号7）
# 数字"9"（类别编号9）
# 这些数字的选择基于：
# 结构复杂度梯度（从最简单的"1"到较复杂的"3"和"9"）
# 包含易混淆对（如"1"和"7"的斜线相似性）
# 覆盖多种笔画特征（直线、曲线、闭合环等）
# 在现实应用中具有典型代表性

@dataclass
class DataHandle:
    module_name: Any  # 传入模块名
    filter_list: List[int] = None

    # 不参与初始化参数
    _x_train: np.ndarray = field(init=False, default=None)
    _y_train: np.ndarray = field(init=False, default=None)
    _x_test: np.ndarray = field(init=False, default=None)
    _y_test: np.ndarray = field(init=False, default=None)
    _y_cls_num: int = field(init=False, default=None)

    def __post_init__(self):
        # 是否满足条件
        if not self.module_name:
            raise ValueError(" DataHandle __init__ module_name 不能为空!")
        if not self.filter_list:
            self.filter_list = [1, 3, 7, 9]
        # 获取数据
        self._x_train, self._x_test, self._y_train, self._y_test = self.load_data()

    def load_data(self) -> Tuple[np.ndarray, ...]:
        """ 加载数据 """
        (X_train, y_train), (X_test, y_test) = self.module_name.load_data()
        # 对数据进行筛选 - 条件
        train_filter = np.isin(y_train, self.filter_list)
        test_filter = np.isin(y_test, self.filter_list)
        # 对数据进行筛选
        X_train = X_train[train_filter]
        X_test = X_test[test_filter]
        y_train = y_train[train_filter]
        y_test = y_test[test_filter]
        # 对特征进行修改形状，并进行归一化
        X_train = X_train.reshape(-1, 28, 28, 1) / 255.
        X_test = X_test.reshape(-1, 28, 28, 1) / 255.
        # 获取标签类别数量
        self._y_cls_num = len(np.unique(y_train))
        # 对标签进行独热
        y_train = self.one_hot_code(y_train, self._y_cls_num)
        y_test = self.one_hot_code(y_test, self._y_cls_num)
        return X_train, X_test, y_train, y_test

    def one_hot_code(self, y, cls_num) -> np.ndarray:
        """ 对指定数据进行独热 """
        return utils.to_categorical(y, cls_num)

    @property
    def return_train_test(self):
        return self._x_train, self._x_test, self._y_train, self._y_test

    @property
    def y_cls_num(self):
        return self._y_cls_num


class MyLeNet5(Model):

    def __init__(self, outputs, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 卷积层
        self.conv = Sequential([
            Conv2D(filters=6, kernel_size=(5, 5), activation='relu'),
            MaxPooling2D(),
            Conv2D(filters=16, kernel_size=(5, 5), activation='relu'),
            MaxPooling2D()
        ])
        # 展平层
        self.flatten = Sequential([
            Flatten()
        ])
        # 全连接层
        self.dense = Sequential([
            Dense(units=120, activation='relu'),
            Dropout(0.2),
            Dense(units=84, activation='relu'),
            Dropout(0.2),
            Dense(units=outputs, activation='softmax')
        ])

    def call(self, inputs, *args, **kwargs):
        x = self.conv(inputs)
        x = self.flatten(x)
        x = self.dense(x)
        return x


@dataclass
class TrainerConfig:

    batch_size : int = 32
    epochs : int = 10

    product_dict = {
        0:1,
        1:3,
        2:7,
        3:9
    }


@dataclass
class ModelTrainer:

    data_handle : DataHandle = None  # 数据处理类
    train_config : TrainerConfig = None  # 训练配置

    _model : MyLeNet5 = field(init=False, default=None)
    _history : keras.src.callbacks.History = field(init=False, default=None)

    def __post_init__(self):
        # 加载数据
        X_train, X_test, y_train, y_test = self.data_handle.return_train_test
        # 实例化模型
        self._model = MyLeNet5(self.data_handle.y_cls_num)
        # 构建模型
        self._model.build(input_shape=X_train.shape)
        # 模型编译
        self._model.compile(
            optimizer=optimizers.Adam(learning_rate=0.001),
            loss=losses.categorical_crossentropy,
            metrics=['acc', metrics.Precision, metrics.Recall, metrics.F1Score]
        )
        self._history = self._model.fit(
            X_train, y_train, batch_size=self.train_config.batch_size,
            epochs=self.train_config.epochs, validation_data=(X_test, y_test)
        )
        """
        模型训练效果较好
         acc: 0.9980 - f1_score: 0.9980 - loss: 0.0055 - precision: 0.9980 - recall: 0.9980 - val_acc: 0.9955 - val_f1_score: 0.9954 - val_loss: 0.0184 - val_precision: 0.9955 - val_recall: 0.9952
         
         模型调优：
            没加失活层：acc: 0.9980 - f1_score: 0.9980 - loss: 0.0055 - precision: 0.9980 - recall: 0.9980 - val_acc: 0.9955 - val_f1_score: 0.9954 - val_loss: 0.0184 - val_precision: 0.9955 - val_recall: 0.9952
            加失活层：- acc: 0.9985 - f1_score: 0.9985 - loss: 0.0048 - precision: 0.9985 - recall: 0.9983 - val_acc: 0.9964 - val_f1_score: 0.9964 - val_loss: 0.0113 - val_precision: 0.9964 - val_recall: 0.9959
            
            加了失活层，Adam 学习率为0.001时：- acc: 0.9985 - f1_score: 0.9985 - loss: 0.0048 - precision: 0.9985 - recall: 0.9983 - val_acc: 0.9964 - val_f1_score: 0.9964 - val_loss: 0.0113 - val_precision: 0.9964 - val_recall: 0.9959
            加了失活层，Adam 学习率为0.005时：- acc: 0.9961 - f1_score: 0.9960 - loss: 0.0148 - precision: 0.9962 - recall: 0.9958 - val_acc: 0.9892 - val_f1_score: 0.9891 - val_loss: 0.0618 - val_precision: 0.9895 - val_recall: 0.9888
            加了失活层，Adam 学习率为0.01时： - acc: 0.9936 - f1_score: 0.9936 - loss: 0.0212 - precision: 0.9970 - recall: 0.9919 - val_acc: 0.9935 - val_f1_score: 0.9935 - val_loss: 0.0482 - val_precision: 0.9952 - val_recall: 0.9916  
        """
        # 保存模型权重
        self._model.save_weights('model.weights.h5')
        self._model.load_weights("model.weights.h5")
        print(self._model.evaluate(X_test, y_test, return_dict=True))


    @property
    def return_model_history(self) -> Tuple[MyLeNet5, Any]:
        return self._model, self._history

def main():
    """ 主程序 """
    data_handle = DataHandle(mnist)  # 实例化数据
    train_config = TrainerConfig()

    # 实例化训练器
    trainer_ = ModelTrainer(data_handle, train_config)
    # 获取训练的模型和训练历史
    model, history = trainer_.return_model_history
    # 打印模型的形状
    model.summary()
    # 获取到模型的训练和验证准确率,绘制图表
    train_acc = history.history['acc']
    val_acc = history.history['val_acc']
    plt.plot(train_acc, label='train acc', c="r")
    plt.plot(val_acc, label='val acc', c="g")
    plt.legend()
    plt.show()

    # 绘制损失图
    plt.plot(history.history['loss'], label='loss')
    plt.title('loss')
    plt.legend()
    plt.show()




if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(str(e))



# 3)	模型编译：
# 使用Adam优化器。
# 损失函数适用于多分类问题。
# 评估指标为准确率。
# 4)	模型训练：
# 使用训练集数据训练模型，指定训练周期（epochs）和批大小（batch_size）。
# 使用测试集数据进行验证。
# 记录训练过程中的关键指标，评估指标为精度（精确率、召回率、F1）和损失值，以便后续分析和调优。
# 训练过程中，可视化展示训练集、验证集的精度、1oss的变化曲线。
# 简要分析模型的训练效果（分析的内容以注释的形式总结出来）
# 5)	性能评估约束：
# 在训练完成后，应在测试集上评估模型的性能，包括准确率、损失等。
# 允许生成和可视化训练过程中性能指标的图表，以便直观地了解模型的训练过程。
# 在测试集上使用保存下来的模型文件评估分类模型的性能，包括准确率、
# 混淆矩阵等指标。
# 6)	模型优化约束：
# 选择一种优化方法对模型进行进一步优化(如调整学习率、使用正则化、剪枝等），并分析优化后的结果。
# 7)	结果可视化：
# 绘制训练和验证过程中的准确率曲线。
# 8)	模型保存：
# 将训练好的模型权重保存到文件中。
# 9)	错误处理和日志记录约束：
# 在数据加载、模型训练、性能评估等过程中，应妥善处理可能出现的错误和异常。
# 记录详细的日志信息，以便在出现问题时进行调试和故障排查。
# 模型调优，考题中要求使用体现不同优化方案的要求，比如超参数方面的，比如说分类精度低的数据集的增强；比如说调整网络结构，比如说调整优化器学习率等；不能只是增加epoch大小，或者调整batch_size就可以，必须采用两种或者两种以上的调优方法进行调优，最终结论要给出综合各种方法调优得出的综合结论
